package com.njuse_czf.kg.web.service.impl;

import com.njuse_czf.kg.web.config.SparkConfig;
import com.njuse_czf.kg.web.service.OptimizeAnswerService;
import io.github.briqt.spark4j.SparkClient;
import io.github.briqt.spark4j.constant.SparkApiVersion;
import io.github.briqt.spark4j.exception.SparkException;
import io.github.briqt.spark4j.model.SparkMessage;
import io.github.briqt.spark4j.model.SparkSyncChatResponse;
import io.github.briqt.spark4j.model.request.SparkRequest;
import io.github.briqt.spark4j.model.response.SparkTextUsage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

@Component
public class SparkOptimizeImpl implements OptimizeAnswerService {

    @Autowired
    private SparkConfig sparkConfig;

    @Override
    public String optimize(String answer){
        List<SparkMessage> messages = new ArrayList<>();

        messages.add(SparkMessage.systemContent("Please help me make the sentences I provide smooth, you just need to return the optimized sentences"));
        messages.add(SparkMessage.userContent(answer));

        return syncFunctionCall(messages);
    }


    private String syncFunctionCall(List<SparkMessage> messages){
        SparkClient client = new SparkClient();
        client.appid = sparkConfig.getAppId();
        client.apiKey = sparkConfig.getApiKey();
        client.apiSecret = sparkConfig.getApiSecret();

        //构造请求
        SparkRequest sparkRequest = SparkRequest.builder()
                .messages(messages)
                .maxTokens(256)
                .temperature(0.2)
                .apiVersion(SparkApiVersion.V3_0)
                .build();

        String response = null;

        try{
            SparkSyncChatResponse chatResponse = client.chatSync(sparkRequest);
            SparkTextUsage textUsage = chatResponse.getTextUsage();

            System.out.println("\n提问tokens：" + textUsage.getPromptTokens()
                    + "，回答tokens：" + textUsage.getCompletionTokens()
                    + "，总消耗tokens：" + textUsage.getTotalTokens());

            response =  chatResponse.getContent();
        }catch (SparkException e){
            System.out.println("发生异常了：" + e.getMessage());
        }

        return response;
    }
}
